﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using RayDen.Library.Components.SystemComponents;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Data;
using Wintellect.PowerCollections;

namespace RayDen.RayEngine.IntersectionEngines
{
    public class LinearBvhAccelleration
    {
        public Node RootNode;
        private int bvhNodesCount;


        public class Node
        {
            public AABB Bound;
            public Node Left, Right;
            public List<BvhPrimitive> Primitives;
            public Axis SplitAxis;
        }

        public class BvhPrimitive
        {
            public AABB Bounds;
            public int Index;
            public uint MortonCode;
        }

        private Point[] sceneVertices;
        public TriangleDataInfo[] triangles;
        private List<ITriangleMesh> meshes;
        private List<BvhPrimitive> bvhPrimitives;


        public LNode[] nodes;
        public int rootOffset;

        public void Init(RayEngineScene scene, Data.TriangleMesh[] geometry)
        {
            var sceneModel = scene;
            meshes = new List<ITriangleMesh>();
            meshes.AddRange(scene.Meshes);
            this.sceneVertices = sceneModel.Vertices;
            this.triangles = scene.Triangles.ToArray();
            bvhPrimitives = this.triangles.Select((item, i) =>
                {
                    var result = new BvhPrimitive() {Bounds = item.WorldBound(sceneVertices), Index = i};
                    result.Bounds.Expand(MathLab.RAY_EPSILON);
                    return result;
                }).ToList();

            this.BuildTree();

            /*
            var tris = new List<TriangleInfo>(triangles.Count());
            tris.AddRange(bvhPrimitives.Select(bvhPrimitive => triangles[bvhPrimitive.Index]));
            this.triangles = tris.ToArray();
            */
            Tracer.TraceLine("Linearize BVH");
            this.LinearizeBVH();
            Tracer.TraceLine("Complete");
            


            Tracer.TraceLine("Tree is builded, {0} nodes total", bvhNodesCount);
        }

        public LNode[] GetData()
        {
            return this.nodes;
        }

        private void LinearizeBVH()
        {
            nodes = new LNode[bvhNodesCount];
            rootOffset = this.BuildArray(RootNode, 0);

            //var znNode = this.nodes.Single(item => item.PrimitiveIndex == 1 && item.IsLeaf);
            //Tracer.TraceLine("ZNODE {0}", znNode.BData);
            Tracer.TraceLine("Root node offset {0}", rootOffset);
        }

        private int BuildArray2(Node node, int offset)
        {
            var nd = node;
            LNode p;
            while (nd != null)
            {
                var firstOffset = offset;
                p = new LNode()
                {
                    Bounds = new BBox(nd.Bound)
                };

                offset = BuildArray(nd.Left, offset + 1);

                if (nd.Primitives != null)
                {
                    p.PrimitiveIndex = (uint)nd.Primitives[0].Index;
                    p.IsLeaf = true;
                }
                else
                {
                    p.Axis = (byte)nd.SplitAxis;
                    p.SkipIndex = (uint)offset;
                }

                nodes[firstOffset] = p;
                nd = nd.Right;
            }

            return offset;
        }
        private int BuildArray(Node node, int offset)
        {
            var nd = node;
            LNode p;
            while (nd != null)
            {
                var firstOffset = offset;
                p = new LNode()
                {
                    Bounds = new BBox(nd.Bound),
                };

                offset = BuildArray(nd.Left, offset + 1);

                if (nd.Primitives != null)
                {
                    p.PrimitiveIndex = (uint)nd.Primitives[0].Index;
                    p.IsLeaf = true;
                }
                else
                {
                    p.Axis = (byte)nd.SplitAxis;
                    p.SkipIndex = (uint)offset;
                }

                nodes[firstOffset] = p;
                nd = nd.Right;
            }

            return offset;
        }
//        private void BuildTree()
//        {

//            var mortonCodes = new List<uint>();
//            AABB bound = bvhPrimitives.Select(prim => prim.Bounds).Aggregate((b1, b2) => b1.Union(b2));
//            sceneBounds = bound;
//            foreach (var bvhPrimitive in bvhPrimitives)
//            {
//                var p = 
//                    //bvhPrimitive.Bounds.Center;
//                    (bound.Center - bvhPrimitive.Bounds.Center) / bound.Size;
//                    //-bound.Center;
//                    //(bvhPrimitive.Bounds.Sum) / bound.Size;
//                    //(bound.Center - bvhPrimitive.Bounds.Center);
//                var mc = Bits.Morton3D
//                    (p);
//                    //(p.x, p.y, p.z);
//                bvhPrimitive.MortonCode = mc;
//                mortonCodes.Add(mc);
//            }
//            mortonCodes.Sort();
//            //mortonCodes.Reverse();

//            bvhPrimitives.Sort((a, b) => a.MortonCode.CompareTo(b.MortonCode));
//            //bvhPrimitives.Reverse();

//            RootNode = new Node() { Bound = bound };
//            bvhNodesCount++;
//            int split = FindSplit(mortonCodes, 0, mortonCodes.Count-1);
//            RootNode.Left = BuildSubTree(mortonCodes, bvhPrimitives, 0, split);
//            RootNode.Right = BuildSubTree(mortonCodes, bvhPrimitives, split + 1, bvhPrimitives.Count-1);
////            RootNode.Bound = RootNode.Left.Bound.Union(RootNode.Right.Bound);
////            RootNode.SplitAxis = (Axis) bound.MaximumExtent();
//        }

        private void BuildTree()
        {
            AABB bound = bvhPrimitives.Select(prim => prim.Bounds).Aggregate((b1, b2) => b1.Union(b2));
            var mortonCodes = new List<uint>();
            foreach (var bvhPrimitive in bvhPrimitives)
            {
                var p = (bound.Center - bvhPrimitive.Bounds.Center) / bound.Size;
                var mc = Bits.Morton3D
                    (p);
                //    (p.x, p.y, p.z);
                bvhPrimitive.MortonCode = mc;
                mortonCodes.Add(mc);
            }
            mortonCodes.Sort();
            bvhPrimitives.Sort((a, b) => a.MortonCode.CompareTo(b.MortonCode));

            RootNode = new Node() { Bound = bound };
            bvhNodesCount++;
            int split = FindSplit(mortonCodes, 0, mortonCodes.Count - 1);
            RootNode.Left = BuildSubTree(mortonCodes, bvhPrimitives, 0, split);
            RootNode.Right = BuildSubTree(mortonCodes, bvhPrimitives, 1 + split, bvhPrimitives.Count - 1);
        }
        private Node BuildSubTree(List<uint> sortedMortonCodes, List<BvhPrimitive> sortedPrimitives, int first,
                           int last)
        {
            if (first == last)
            {
                bvhNodesCount++;
                return new Node()
                {
                    Bound = sortedPrimitives[first].Bounds,
                    Primitives = new List<BvhPrimitive>() { sortedPrimitives[first] }
                };
            }
            AABB bound = sortedPrimitives.GetRange(first, last - first).Select(prim => prim.Bounds).Aggregate((b1, b2) => b1.Union(b2));


            int split = FindSplit(sortedMortonCodes, first, last);


            var childA = BuildSubTree(sortedMortonCodes, sortedPrimitives,
                                    first, split);
            var childB = BuildSubTree(sortedMortonCodes, sortedPrimitives,
                                             split + 1, last);
            bvhNodesCount++;
            return new Node() { Bound = childA.Bound.Union(childA.Bound).Union(childB.Bound), Left = childA, Right = childB };
        }

        private int FindSplit(List<uint> sortedMortonCodes, int first, int last)
        {
            uint firstCode = sortedMortonCodes[first];
            uint lastCode = sortedMortonCodes[last];

            if (firstCode == lastCode)
                return (first + last) >> 1;

            // Calculate the number of highest bits that are the same
            // for all objects, using the count-leading-zeros intrinsic.

            int commonPrefix = Bits.lzc(firstCode ^ lastCode);

            // Use binary search to find where the next bit differs.
            // Specifically, we are looking for the highest object that
            // shares more than commonPrefix bits with the first one.

            int split = first; // initial guess
            int step = last - first;

            do
            {
                step = (step + 1) >> 1; // exponential decrease
                int newSplit = split + step; // proposed new position

                if (newSplit < last)
                {
                    uint splitCode = sortedMortonCodes[newSplit];
                    int splitPrefix = Bits.lzc(firstCode ^ splitCode);
                    if (splitPrefix > commonPrefix)
                        split = newSplit; // accept proposal
                }
            }
            while (step > 1);

            return split;
        }
    }
}
